import torch

input_ids = torch.tensor([2, 3, 5, 1])  # 定义输入词元ID张量

vocab_size = 6  # 词汇表大小
output_dim = 3  # 输出维度

torch.manual_seed(123)  # 设置随机种子
embedding_layer = torch.nn.Embedding(vocab_size, output_dim)  # 实例化嵌入层
print(embedding_layer.weight)  # 打印嵌入层的权重矩阵
'''
tensor([[ 0.3374, -0.1778, -0.1690],
        [ 0.9178,  1.5810,  1.3010],
        [ 1.2753, -0.2010, -0.1606],
        [-0.4015,  0.9666, -1.1481],
        [-1.1589,  0.3255, -0.6315],
        [-2.8400, -0.7849, -1.4096]], requires_grad=True)
'''
print(embedding_layer(input_ids))  # 获取所有输入词元ID的嵌入向量
'''
tensor([[ 1.2753, -0.2010, -0.1606],
        [-0.4015,  0.9666, -1.1481],
        [-2.8400, -0.7849, -1.4096],
        [ 0.9178,  1.5810,  1.3010]], grad_fn=<EmbeddingBackward0>)
'''

